from matchteris.env import Assemble_Suction
import mujoco
import numpy as np
from MoveBlock.get_env_data import GetEnvData
from MoveBlock.pos_trainsport import PosTrains

posThreshold = 5e-4         # 位置误差阈值
angleThreshold = 5e-4       # 角度误差阈值

distance_off = np.array([1.5e-3, 1e-3, 0])   #位置补偿
angle_off = 0.05                             #角度补偿


class ArmControl:
    '''
        用于控制机械臂移动
    '''
    overTimes = 0           # 超时次数，当机械臂到达指定位置时清零
    spinOverTime = 0        # 旋转超时次数，当机械臂旋转到指定角度时清零
    maxOverTimes = 1000     # 最大超时时间，当达到最大超时时间时认为机械臂已经到达指定位置
    maxSpinOverTimes = 100  # 最大旋转超时时间，当达到最大旋转超时时间时认为机械臂已经旋转到指定角度 

    @classmethod
    def isAttach(cls, env:Assemble_Suction) -> bool:
        '''
            判断机械臂是否吸附方块
        '''
        return env.data.ctrl[-1] == 1

    @classmethod
    def isArarive(cls, env:Assemble_Suction, destination:np.ndarray, threshold:float) -> bool:
        '''
            判断机械臂是否到达指定位置
            destination: (x, y, z)
            threshold: 误差阈值
        '''
        current_pos = GetEnvData.GetSuckerPosition(env)
        off = np.max(np.abs(current_pos/destination - 1))
        if off < threshold or cls.overTimes > cls.maxOverTimes:
            cls.overTimes = 0
            return True
        else:
            cls.overTimes += 1
            return False
    
    @classmethod
    def isRotateArarive(cls, env:Assemble_Suction, destination:float, threshold:float) -> bool:
        '''
            判断机械臂是否到达指定角度
            destination: 角度
            threshold: 误差阈值
        '''
        current_angle = GetEnvData.GetSuckerAngle(env)
        off = np.abs(current_angle/destination - 1)
        if off < threshold or cls.spinOverTime > cls.maxSpinOverTimes:
            cls.spinOverTime = 0
            return True
        else:
            cls.spinOverTime += 1
            return False

    @classmethod
    def Suck(cls, env:Assemble_Suction, destination:np.ndarray) -> None:
        '''
            机械臂吸取方块
            destination: (x, y, z), 吸取位置
        '''
        for _ in range(100):
            env.data.mocap_pos[env.mocap_id] = destination
            env.data.ctrl[-1] = 1
            env.mocap_ctrl()

    @classmethod
    def Release(cls, env:Assemble_Suction, destination:np.ndarray) -> None:
        '''
            机械臂释放方块
            destination: (x, y, z), 释放位置
        '''
        for _ in range(100):
            env.data.mocap_pos[env.mocap_id] = destination
            env.data.ctrl[-1] = 0
            env.mocap_ctrl()

    @classmethod
    def MoveOnWay(cls, env:Assemble_Suction, way:int, distance:float, step:int, threshold:float) -> None:
        '''
            机械臂前后移动
            way: 0为前后移动(x轴), 1为左右移动(y轴), 2为上下移动(z轴)
            distance: 距离
            step: 步数， 数值越大移动越慢，但是更平滑，会根据是否吸附方块调整，吸附方块时step*=2，不吸附方块时step/=10
            threshold: 误差阈值
        '''
        # 设置超时时间
        maxOverTimes = cls.maxOverTimes
        cls.maxOverTimes += step

        # 根据吸附情况调整步数
        if cls.isAttach(env):
            step *= 2
        else:
            step /= 10

        # 计算与目标位置的偏移量
        off = np.zeros(3)
        off[way] = distance
        # 计算目标位置
        destination = GetEnvData.GetSuckerPosition(env).copy() + off
        # 每步移动的距离
        delta = distance/step
        count = 0
        # 移动到目标位置
        while not cls.isArarive(env, destination, threshold):
            if count < step:
                env.data.mocap_pos[env.mocap_id][way] += delta
                count += 1
            env.mocap_ctrl()
        # 保持在目标位置
        env.data.mocap_pos[env.mocap_id] = destination
        env.mocap_ctrl()
        while not cls.isArarive(env, destination, threshold):
            env.mocap_ctrl()

        # 重置超时时间
        cls.maxOverTimes = maxOverTimes

    @classmethod
    def StepMove(cls, env:Assemble_Suction, destination:np.ndarray, step:int,
                  threshold:float, distance_off:np.ndarray[float]) -> None:
        '''
            机械臂移动
            distance: 距离
            step: 步数, 数值越大移动越慢，但是更平滑，会根据是否吸附方块调整，吸附方块时step*=2，不吸附方块时step/=10
            threshold: 误差阈值
            distance_off: 距离补偿
        '''
        # 设置超时时间
        maxOverTimes = cls.maxOverTimes
        cls.maxOverTimes += step

        # 根据吸附情况调整步数
        if cls.isAttach(env):
            step *= 2
        else:
            step /= 10

        # 补偿距离
        destination += distance_off

        # 计算与目标位置的偏移量
        current_pos = GetEnvData.GetSuckerPosition(env).copy()
        off = (destination - current_pos)/step
        count = 0
        # 移动到目标位置
        while not cls.isArarive(env, destination, threshold):
            if count < step:
                count += 1
                env.data.mocap_pos[env.mocap_id] += off
            env.mocap_ctrl()
        # 保持在目标位置
        env.data.mocap_pos[env.mocap_id] = destination
        env.mocap_ctrl()
        while not cls.isArarive(env, destination, threshold):
            env.mocap_ctrl()

        # 重置超时时间
        cls.maxOverTimes = maxOverTimes

    @classmethod
    def SetSpinAngle(cls, env:Assemble_Suction, degree:float) -> None:
        '''
            设置机械臂旋转角度
            degree: 角度
        '''
        degree = np.clip(degree, -180, 180)
        theta = np.deg2rad(degree)
        rotate_xmat = np.array([[np.cos(theta),-np.sin(theta),0],[np.sin(theta),np.cos(theta),0],[0,0,1]])
        mocap_quat = env.data.mocap_quat[env.mocap_id]
        mocap_xmat = np.zeros(9)
        mujoco.mju_quat2Mat(mocap_xmat,mocap_quat)
        mocap_xmat = mocap_xmat.reshape(3,3)
        mocap_xmat = (rotate_xmat@mocap_xmat).flatten()
        mujoco.mju_mat2Quat(mocap_quat,mocap_xmat)
        env.data.mocap_quat[env.mocap_id] = mocap_quat
    
    @classmethod
    def StepSpin(cls, env:Assemble_Suction, degree:float, threshold:float, angle_off:float=angle_off) -> None:
        '''
            机械臂旋转指定角度
            degree: 角度
            threshold: 误差阈值
            angle_off: 角度补偿
        '''
        # 设置旋转步数
        step = 50

        # 补偿角度
        degree *= 1 + angle_off

        # 计算目标角度，由于角度在仿真环境为负数，因此减去目标值。
        destination = GetEnvData.GetSuckerAngle(env) - degree

        # 计算每步旋转的角度
        delta = degree/step
        count = 0

        # 旋转到指定角度
        while count < step:
            count += 1
            ArmControl.SetSpinAngle(env, delta)
            for _ in range(10):
                env.mocap_ctrl()

        # 保持在目标角度
        while not cls.isRotateArarive(env, destination, threshold):
            env.rotate_mocap(GetEnvData.GetSuckerAngle(env) - destination)

    @classmethod
    def MoveBlock(cls, env:Assemble_Suction, block_name:str, destination:np.ndarray, spin:float=0, 
                  posThreshold:float=posThreshold, angleThreshold:float=angleThreshold, step:int=1000,
                  distance_off:np.ndarray[float]=distance_off, angle_off:float=angle_off) -> None:
        '''
            机械臂移动到指定方块上方，吸取方块，然后移动到指定位置, 可以指定旋转角度
            block_name: color + index + "/", 如"orange0/"
            destination: (x, y, z)
            spin: 旋转角度
            posThreshold: 位置误差阈值
            angleThreshold: 角度误差阈值
            step: 步数， 数值越大移动越慢，但是更平滑，会根据是否吸附方块调整，吸附方块时step*=2，不吸附方块时step/=10
            distance_off: 距离补偿
            angle_off: 角度补偿
        '''
        # 获取方块位置
        block_pos = GetEnvData.GetBlockPosition(env, block_name)
        # 转换方块坐标，将中心坐标转换到方格上，以便吸附
        block_pos = PosTrains.TrainsPortBlockPos(block_name, block_pos)

        # 机械臂移动到方块上方
        ArmControl.StepMove(env, block_pos, step, posThreshold, distance_off)
        # 机械臂吸取方块
        ArmControl.Suck(env, block_pos)
        # 举起一段高度，防止碰撞
        ArmControl.MoveOnWay(env, 2, PosTrains.tetris_height*10 , step, posThreshold)

        # 机械臂移动到指定位置上方
        ArmControl.StepMove(env, destination + np.array([0, 0, 10 * PosTrains.tetris_height]),
                             step, posThreshold, distance_off)

        # 机械臂旋转指定角度
        if spin != 0:
            ArmControl.StepSpin(env, spin, angleThreshold, angle_off)

        # 机械臂移动到指定位置
        ArmControl.StepMove(env, destination, step, posThreshold, distance_off)
        # 机械臂释放方块
        ArmControl.Release(env, destination)

        # 重新举高防止碰撞
        ArmControl.MoveOnWay(env, 2, PosTrains.tetris_height*10, step, posThreshold)

        # 旋转角度归0
        if spin!= 0:
            ArmControl.StepSpin(env, -spin, angleThreshold, angle_off)

class BlockControl:
    '''
        用于移动方块到盘子中的指定位置
    '''

    @classmethod
    def InitEnv(cls, env:Assemble_Suction) -> None:
        '''
            初始化机械臂, 生成方块
        '''
        env.generate_blocks()
        for _ in range(500):
            env.mocap_ctrl()

    @classmethod
    def MoveBlock(cls, env:Assemble_Suction, block_name:str, destination_x:int, destination_y:int, spin:float=0,
                posThreshold:float=posThreshold, angleThreshold:float=angleThreshold, step:int=1000,
                distance_off:np.ndarray[float]=distance_off, angle_off:float=angle_off) -> None:
        '''
            机械臂移动到指定方块上方，吸取方块，然后移动到盘子中的指定位置, 可以指定旋转角度
            block_name: color + index + "/", 如"orange0/"
            destination_x: 盘子x坐标
            destination_y: 盘子y坐标
            spin: 旋转角度
            posThreshold: 位置误差阈值
            angleThreshold: 角度误差阈值
            step: 步数， 数值越大移动越慢，但是更平滑，会根据是否吸附方块调整，吸附方块时step*=2，不吸附方块时step/=10
            distance_off: 距离补偿
            angle_off: 角度补偿
        '''
        # 获取目标位置
        destination = PosTrains.TrainSport(destination_x, destination_y)

        # 机械臂移动到指定方块上方，吸取方块，然后移动到指定位置
        ArmControl.MoveBlock(env, block_name, destination, spin, posThreshold, 
                             angleThreshold, step, distance_off, angle_off)
        
    @classmethod
    def PosError(cls, src:np.ndarray, dst:np.ndarray) -> float:
        '''
            计算两个位置的误差
            src: 源位置
            dst: 目标位置
            return: 误差
        '''
        return dst/src - 1 


if __name__ == '__main__':
    env = Assemble_Suction()
    env.generate_blocks()
    BlockControl.InitEnv(env)
    target = PosTrains.TrainSport(2, 2)
    taeget2 = PosTrains.TrainSport(4, 4)
    ArmControl.StepSpin(env, 90, angleThreshold)
    ArmControl.StepSpin(env, -90, angleThreshold)

    BlockControl.MoveBlock(env, "blue0/", 2, 2, step=1000, spin=0)
    BlockControl.MoveBlock(env, "red0/", 4, 4, step=1000, spin=90)
    block_pos = PosTrains.TrainsPortBlockPos("blue0/", GetEnvData.GetBlockPosition(env, "blue0/"))
    block_pos2 = PosTrains.TrainsPortBlockPos("red0/", GetEnvData.GetBlockPosition(env, "red0/"))
    print("\n最终误差1: ",BlockControl.PosError(block_pos, target))
    print("\n最终误差2: ",BlockControl.PosError(block_pos2, taeget2))